package http_proxy_middleware

import (
	"fmt"
	"gitee.com/zhoulvvv/my_go_gateway/dao"
	"gitee.com/zhoulvvv/my_go_gateway/middleware"
	"gitee.com/zhoulvvv/my_go_gateway/public"
	"github.com/gin-gonic/gin"
	"github.com/pkg/errors"
)

// HTTP 服务端和客户端限流中间件
func HTTPFlowLimitModeMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		serviceInterface, ok := c.Get("service")
		if !ok {
			middleware.ResponseError(c, 2001, errors.New("【服务不存在】service not found"))
			c.Abort()
			return
		}
		service := serviceInterface.(*dao.ServiceDetail)
		serverQps := float64(service.AccessControl.ServiceFlowLimit)
		// 如果服务端限流QPS是0就表示无限制
		if serverQps > 0 {
			// 服务端限流
			serverLimiter, err := public.FlowLimiterHandler.GetLimiter(
				public.FlowServicePrefix+service.Info.ServiceName,
				serverQps)
			if err != nil {
				middleware.ResponseError(c, 5001, errors.New("获取服务端限流器失败："+err.Error()))
				c.Abort()
				return
			}
			if !serverLimiter.Allow() {
				middleware.ResponseError(c, 5001, errors.New(fmt.Sprintf("服务端限流中，每秒最大QPS：%v", serverQps)))
				c.Abort()
				return
			}
		}

		clientQps := float64(service.AccessControl.ClientIPFlowLimit)
		// 如果客户端限流QPS是0就表示无限制
		if clientQps > 0 {
			// 客户端限流，客户端限流是基于IP进行限流
			clientLimiter, err := public.FlowLimiterHandler.GetLimiter(
				public.FlowServicePrefix+service.Info.ServiceName+"_"+c.ClientIP(),
				clientQps)
			if err != nil {
				middleware.ResponseError(c, 5001, errors.New("获取客户端限流器失败："+err.Error()))
				c.Abort()
				return
			}
			if !clientLimiter.Allow() {
				middleware.ResponseError(c, 5001, errors.New(fmt.Sprintf("客户端限流中，每秒最大QPS：%v，IP：%s", clientQps, c.ClientIP())))
				c.Abort()
				return
			}
		}
		c.Next()
	}
}
